﻿using FFmpeg.AutoGen;
using NLog;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.Threading;
using System.Threading.Tasks;

namespace VideoWriter
{
    /// <summary>
    /// Mp4文件写入类
    /// </summary>
    public unsafe class Mp4VideoWriter : IDisposable
    {
        #region 字段

        /// <summary>
        /// 正在写
        /// </summary>
        private ManualResetEventSlim isWritted = new ManualResetEventSlim(false);
        private bool canWrite;

        private Logger logger = LogManager.GetCurrentClassLogger();


        /// <summary>
        /// 音频PTS
        /// </summary>
        private long audioPts;

        private object lock_obj = new object();

        /// <summary>
        /// PCM音频数据缓冲
        /// </summary>
        private List<byte> pcmBuffer = new List<byte>();

        private readonly Size _frameSize;
        private readonly int _linesizeU;
        private readonly int _linesizeV;
        private readonly int _linesizeY;
        private readonly AVCodec* _pVideoCodec, _pAudioCodec;
        private readonly AVStream* _pVideoStream, _pAudioStream;
        private readonly AVCodecContext* _pVideoCodecContext, _pAudioCodecContext;
        private readonly AVFormatContext* _pFormatContext;
        private readonly AVOutputFormat* _pOutputFormat;
        private readonly SwrContext* _pSwrContext;
        private readonly AVFrame* _pInAudioFrame;
        private readonly AVFrame* _pOutAudioFrame;
        private readonly AVFrame* _pInVideoFrame;
        private readonly AVFrame* _pOutVideoFrame;

        private readonly int _uSize;
        private readonly int _ySize;

        private readonly int _fps;

        private VideoFrameConverter videoFrameConverter;


        /// <summary>
        /// 已写入的帧数
        /// </summary>
        private long framePts;

        private long pktPts;

        /// <summary>
        /// 已写入的音频采样点的个数
        /// </summary>
        private long samplesCount;

        #endregion

        #region 公共方法

        /// <summary>
        /// 创建视频写入对象
        /// </summary>
        /// <param name="fileName"></param>
        /// <param name="fps"></param>
        /// <param name="frameSize"></param>
        public Mp4VideoWriter(string fileName, int fps, Size frameSize)
        {

            #region Init

            _fps = fps;
            var sourcePixelFormat = AVPixelFormat.AV_PIX_FMT_BGR24;
            var destinationPixelFormat = AVPixelFormat.AV_PIX_FMT_YUV420P;

            videoFrameConverter = new VideoFrameConverter(frameSize, sourcePixelFormat, frameSize, destinationPixelFormat);

            _frameSize = frameSize;
            fixed (AVFormatContext** p = &_pFormatContext)
                ffmpeg.avformat_alloc_output_context2(p, null, "mp4", fileName);
            if (_pFormatContext == null) throw new InvalidOperationException("alloc output context failed.");

            _pOutputFormat = _pFormatContext->oformat;

            _pOutputFormat->video_codec = AVCodecID.AV_CODEC_ID_H264;
            _pOutputFormat->audio_codec = AVCodecID.AV_CODEC_ID_AAC;

            _pVideoCodec = ffmpeg.avcodec_find_encoder(AVCodecID.AV_CODEC_ID_H264);
            if (_pVideoCodec == null) throw new InvalidOperationException("Video Codec not found.");

            _pAudioCodec = ffmpeg.avcodec_find_encoder(AVCodecID.AV_CODEC_ID_AAC);
            if (_pAudioCodec == null) throw new InvalidOperationException("Audio Codec not found.");

            /* 添加视频输出流 */
            {
                var st = ffmpeg.avformat_new_stream(_pFormatContext, _pVideoCodec);
                if (st == null) throw new InvalidOperationException("creat stream failed.");
                st->id = (int)_pFormatContext->nb_streams - 1;
                _pVideoCodecContext = st->codec;
                st->time_base = new AVRational { num = 1, den = 90000 };
                _pVideoCodecContext->time_base = new AVRational { num = 1, den = _fps };
                _pVideoCodecContext->width = frameSize.Width;
                _pVideoCodecContext->height = frameSize.Height;
                _pVideoCodecContext->gop_size = 12;
                _pVideoCodecContext->pix_fmt = AVPixelFormat.AV_PIX_FMT_YUV420P;
                if ((_pFormatContext->oformat->flags & ffmpeg.AVFMT_GLOBALHEADER) != 0)
                {
                    _pVideoCodecContext->flags |= ffmpeg.AV_CODEC_FLAG_GLOBAL_HEADER;
                }

                _pVideoStream = st;
            }

            /* 添加音频输出流 */
            {
                var st = ffmpeg.avformat_new_stream(_pFormatContext, _pAudioCodec);
                if (st == null) throw new InvalidOperationException("creat stream failed.");
                st->id = (int)_pFormatContext->nb_streams - 1;
                _pAudioCodecContext = st->codec;
                _pAudioCodecContext->codec_type = AVMediaType.AVMEDIA_TYPE_AUDIO;
                _pAudioCodecContext->sample_fmt = AVSampleFormat.AV_SAMPLE_FMT_FLTP;
                _pAudioCodecContext->bit_rate = 64000;
                _pAudioCodecContext->sample_rate = 44100;
                _pAudioCodecContext->channel_layout = ffmpeg.AV_CH_LAYOUT_MONO;
                _pAudioCodec->channel_layouts = &_pAudioCodecContext->channel_layout;
                _pAudioCodecContext->channels = ffmpeg.av_get_channel_layout_nb_channels(_pAudioCodecContext->channel_layout);
                st->time_base = new AVRational { num = 1, den = _pAudioCodecContext->sample_rate };
                //_pAudioCodecContext->time_base = st->time_base;

                if ((_pFormatContext->oformat->flags & ffmpeg.AVFMT_GLOBALHEADER) != 0)
                {
                    _pAudioCodecContext->flags |= ffmpeg.AV_CODEC_FLAG_GLOBAL_HEADER;
                }

                _pAudioStream = st;
            }

            /*打开视频编码器*/
            {
                AVCodecContext* c = _pVideoStream->codec;
                ffmpeg.avcodec_open2(c, _pVideoCodec, null).ThrowExceptionIfError();
                _pInVideoFrame = alloc_video_frame(AVPixelFormat.AV_PIX_FMT_BGR24, _frameSize.Width, _frameSize.Height);
                _pOutVideoFrame = alloc_video_frame(AVPixelFormat.AV_PIX_FMT_YUV420P, _frameSize.Width, _frameSize.Height);
            }

            #region InitAudio
            {
                int nb_samples;
                //打开音频编码器
                ffmpeg.avcodec_open2(_pAudioCodecContext, _pAudioCodec, null).ThrowExceptionIfError();

                if ((_pAudioCodecContext->codec->capabilities & ffmpeg.AV_CODEC_CAP_VARIABLE_FRAME_SIZE) != 0)
                    nb_samples = 10000;
                else
                    nb_samples = _pAudioCodecContext->frame_size;

                //创建全局的音频帧
                _pInAudioFrame = alloc_audio_frame(AVSampleFormat.AV_SAMPLE_FMT_S16P, _pAudioCodecContext->channel_layout, _pAudioCodecContext->sample_rate, nb_samples);
                _pOutAudioFrame = alloc_audio_frame(_pAudioCodecContext->sample_fmt, _pAudioCodecContext->channel_layout, _pAudioCodecContext->sample_rate, nb_samples);

                //拷贝音频流参数到编码器
                ffmpeg.avcodec_parameters_from_context(_pAudioStream->codecpar, _pAudioCodecContext).ThrowExceptionIfError();

                //配置音频格式转换 从 AV_SAMPLE_FMT_S16P 输入转 AV_SAMPLE_FMT_FLTP
                _pSwrContext = ffmpeg.swr_alloc();
                if (_pSwrContext == null) throw new InvalidOperationException("Creat SwrContext failed.");
                ffmpeg.av_opt_set_int(_pSwrContext, "in_channel_count", _pAudioCodecContext->channels, 0);
                ffmpeg.av_opt_set_int(_pSwrContext, "in_sample_rate", _pAudioCodecContext->sample_rate, 0);
                ffmpeg.av_opt_set_sample_fmt(_pSwrContext, "in_sample_fmt", AVSampleFormat.AV_SAMPLE_FMT_S16P, 0);
                ffmpeg.av_opt_set_int(_pSwrContext, "out_channel_count", _pAudioCodecContext->channels, 0);
                ffmpeg.av_opt_set_int(_pSwrContext, "out_sample_rate", _pAudioCodecContext->sample_rate, 0);
                ffmpeg.av_opt_set_sample_fmt(_pSwrContext, "out_sample_fmt", _pAudioCodecContext->sample_fmt, 0);
                ffmpeg.swr_init(_pSwrContext).ThrowExceptionIfError();
            }

            #endregion

            ffmpeg.av_dump_format(_pFormatContext, 0, fileName, 1);

            /* open the output file, if needed */
            if ((_pOutputFormat->flags & ffmpeg.AVFMT_NOFILE) == 0)
            {
                ffmpeg.avio_open(&_pFormatContext->pb, fileName, ffmpeg.AVIO_FLAG_WRITE).ThrowExceptionIfError();
            }
            {
                //AVDictionary* opt;
                //ffmpeg.av_dict_set_int(&opt, "video_track_timescale", _fps, 0);
                ffmpeg.avformat_write_header(_pFormatContext, null).ThrowExceptionIfError();
                isWritted.Set();
                _linesizeY = frameSize.Width;
                _linesizeU = frameSize.Width / 2;
                _linesizeV = frameSize.Width / 2;

                _ySize = _linesizeY * frameSize.Height;
                _uSize = _linesizeU * frameSize.Height / 2;
            }

            canWrite = true;
            #endregion

        }

        /// <summary>
        /// 编码视频帧
        /// </summary>
        /// <param name="bitmapData">BGR图像数据</param>
        /// <param name="timeSpan">距视频创建的时间</param>
        public void WriteFrame(byte[] bitmapData, TimeSpan? timeSpan = null)
        {
            if (!canWrite) return;
            isWritted.Reset();
            if (timeSpan != null)
            {
                framePts = (long)(timeSpan.Value.TotalSeconds * _fps);
            }
            long pkt_pts = ffmpeg.av_rescale_q(framePts, _pVideoCodecContext->time_base, _pVideoStream->time_base);
            if (pkt_pts <= pktPts && pkt_pts > 0)
                return;
            else
                pktPts = pkt_pts;

            fixed (byte* pBitmapData = bitmapData)
            {
                var data = new byte_ptrArray8 { [0] = pBitmapData };
                var linesize = new int_array8 { [0] = bitmapData.Length / _frameSize.Height };

                _pInVideoFrame->data = data;
                _pInVideoFrame->linesize = linesize;
                _pInVideoFrame->height = _frameSize.Height;
                videoFrameConverter.Convert(_pInVideoFrame, _pOutVideoFrame);

                _pOutVideoFrame->pts = framePts;
                if (_pOutVideoFrame->format != (int)_pVideoCodecContext->pix_fmt) throw new ArgumentException("Invalid pixel format.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->width != _frameSize.Width) throw new ArgumentException("Invalid width.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->height != _frameSize.Height) throw new ArgumentException("Invalid height.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->linesize[0] != _linesizeY) throw new ArgumentException("Invalid Y linesize.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->linesize[1] != _linesizeU) throw new ArgumentException("Invalid U linesize.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->linesize[2] != _linesizeV) throw new ArgumentException("Invalid V linesize.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->data[1] - _pOutVideoFrame->data[0] != _ySize) throw new ArgumentException("Invalid Y data size.", nameof(_pOutVideoFrame));
                if (_pOutVideoFrame->data[2] - _pOutVideoFrame->data[1] != _uSize) throw new ArgumentException("Invalid U data size.", nameof(_pOutVideoFrame));
                var pPacket = ffmpeg.av_packet_alloc();

                try
                {
                    int error;
                    do
                    {
                        ffmpeg.avcodec_send_frame(_pVideoCodecContext, _pOutVideoFrame).ThrowExceptionIfError();
                        error = ffmpeg.avcodec_receive_packet(_pVideoCodecContext, pPacket);
                    } while (error == ffmpeg.AVERROR(ffmpeg.EAGAIN));
                    error.ThrowExceptionIfError();

                    pPacket->stream_index = _pVideoStream->index;
                    pPacket->pts = pktPts;
                    //必须要设置dts 否则文件属性中帧率不对
                    pPacket->dts = pPacket->pts;
                    ffmpeg.av_interleaved_write_frame(_pFormatContext, pPacket).ThrowExceptionIfError();
                }
                finally
                {
                    ffmpeg.av_packet_unref(pPacket);
                }
                if (timeSpan == null)
                    framePts++;
            }
            isWritted.Set();
        }


        /// <summary>
        /// 写入PCM音频数据
        /// </summary>
        /// <param name="pcmData">音频采样数据 S16格式</param>
        public void WritePcmAudioFrame(byte[] pcmData)
        {
            //lock (lock_obj)
            {
                pcmBuffer.AddRange(pcmData);
                int bufferSize = ffmpeg.av_get_bytes_per_sample(AVSampleFormat.AV_SAMPLE_FMT_S16P) * _pInAudioFrame->nb_samples * _pAudioCodecContext->channels;
                byte[] inBuffer = new byte[bufferSize];
                while (pcmBuffer.Count >= bufferSize)
                {
                    pcmBuffer.CopyTo(0, inBuffer, 0, bufferSize);
                    pcmBuffer.RemoveRange(0, bufferSize);
                    fixed (byte* data = inBuffer)
                    {
                        _pInAudioFrame->data[0] = data;
                    }
                    EncodeAudio(AVSampleFormat.AV_SAMPLE_FMT_S16P);
                }
            }
        }

        /// <summary>
        /// 写入PCM音频数据
        /// </summary>
        /// <param name="pcmData">音频采样数据 float格式</param>
        public void WriteFloatPcmAudioFrame(byte[] pcmData)
        {
            pcmBuffer.AddRange(pcmData);
            int bufferSize = ffmpeg.av_get_bytes_per_sample(AVSampleFormat.AV_SAMPLE_FMT_FLTP) * _pOutAudioFrame->nb_samples * _pAudioCodecContext->channels;
            byte[] inBuffer = new byte[bufferSize];
            while (pcmBuffer.Count >= bufferSize)
            {
                pcmBuffer.CopyTo(0, inBuffer, 0, bufferSize);
                pcmBuffer.RemoveRange(0, bufferSize);
                fixed (byte* data = inBuffer)
                {
                    _pOutAudioFrame->data[0] = data;
                }
                EncodeAudio(AVSampleFormat.AV_SAMPLE_FMT_FLTP);
            }
        }

        /// <summary>
        /// 写入AAC音频数据
        /// </summary>
        /// <param name="aacData">aac数据</param>
        public void WriteAACAudioFrame(IntPtr aacData, int size)
        {
            size -= 7;
            byte* data = (byte*)aacData;
            data += 7;
            var pPacket = ffmpeg.av_packet_alloc();
            try
            {
                pPacket->size = size;
                pPacket->duration = _pOutAudioFrame->nb_samples;
                pPacket->data = data;
                pPacket->pts = ffmpeg.av_rescale_q(audioPts, _pAudioCodecContext->time_base, _pAudioStream->time_base);
                //必须要设置dts 否则文件属性中帧率不对
                pPacket->dts = pPacket->pts;
                pPacket->stream_index = _pAudioStream->index;
                ffmpeg.av_interleaved_write_frame(_pFormatContext, pPacket).ThrowExceptionIfError();
                audioPts += _pOutAudioFrame->nb_samples;
            }
            finally
            {
                ffmpeg.av_packet_unref(pPacket);
            }
        }

        /// <summary>
        /// 释放
        /// </summary>
        public void Dispose()
        {
            isWritted.Wait();
            canWrite = false;
            logger.Info("Mp4VideoWriter -> Dispose() begin");
            lock (lock_obj)
            {
                ffmpeg.av_write_trailer(_pFormatContext).ThrowExceptionIfError();
                ffmpeg.avcodec_close(_pVideoStream->codec).ThrowExceptionIfError();
                ffmpeg.avcodec_close(_pAudioStream->codec).ThrowExceptionIfError();

                fixed (AVFrame** p = &_pInAudioFrame)
                    ffmpeg.av_frame_free(p);
                fixed (AVFrame** p = &_pOutAudioFrame)
                    ffmpeg.av_frame_free(p);

                fixed (SwrContext** p = &_pSwrContext)
                    ffmpeg.swr_free(p);
                if ((_pOutputFormat->flags & ffmpeg.AVFMT_NOFILE) == 0)
                {
                    ffmpeg.avio_close(_pFormatContext->pb).ThrowExceptionIfError();
                }
                ffmpeg.avformat_free_context(_pFormatContext);
                videoFrameConverter.Dispose();
            }
            logger.Info("Mp4VideoWriter -> Dispose() end");
        }

        #endregion


        #region 内部方法

        /// <summary>
        /// 编码音频帧
        /// </summary>
        /// <param name="sampleFormat">音频采样数据类型</param>
        private void EncodeAudio(AVSampleFormat sampleFormat = AVSampleFormat.AV_SAMPLE_FMT_FLTP)
        {
            if (!canWrite) return;
            if (_pInAudioFrame != null)
            {
                var pPacket = ffmpeg.av_packet_alloc();
                try
                {
                    long dst_nb_samples = ffmpeg.av_rescale_rnd(ffmpeg.swr_get_delay(_pSwrContext, _pAudioCodecContext->sample_rate) + _pInAudioFrame->nb_samples,
                        _pAudioCodecContext->sample_rate, _pAudioCodecContext->sample_rate, AVRounding.AV_ROUND_UP);
                    ffmpeg.av_frame_make_writable(_pOutAudioFrame).ThrowExceptionIfError();
                    if (sampleFormat != AVSampleFormat.AV_SAMPLE_FMT_FLTP)
                    {
                        //转换音频格式
                        fixed (byte** inData = _pInAudioFrame->data.ToArray())
                        fixed (byte** outData = _pOutAudioFrame->data.ToArray())
                        {
                            ffmpeg.swr_convert(_pSwrContext, outData, _pOutAudioFrame->nb_samples, inData, _pInAudioFrame->nb_samples).ThrowExceptionIfError();
                        }
                    }

                    _pOutAudioFrame->pts = ffmpeg.av_rescale_q(samplesCount, _pAudioStream->time_base, _pAudioCodecContext->time_base);
                    samplesCount += _pInAudioFrame->nb_samples;
                    int error;
                    do
                    {
                        ffmpeg.avcodec_send_frame(_pAudioCodecContext, _pOutAudioFrame).ThrowExceptionIfError();
                        error = ffmpeg.avcodec_receive_packet(_pAudioCodecContext, pPacket);
                    } while (error == ffmpeg.AVERROR(ffmpeg.EAGAIN));
                    error.ThrowExceptionIfError();

                    pPacket->pts = ffmpeg.av_rescale_q(_pOutAudioFrame->pts, _pAudioCodecContext->time_base, _pAudioStream->time_base);
                    //必须要设置dts 否则文件属性中帧率不对
                    pPacket->dts = pPacket->pts;
                    pPacket->stream_index = _pAudioStream->index;
                    ffmpeg.av_interleaved_write_frame(_pFormatContext, pPacket).ThrowExceptionIfError();
                }
                finally
                {
                    ffmpeg.av_packet_unref(pPacket);
                }
            }
        }

        /// <summary>
        /// 创建音频帧
        /// </summary>
        /// <param name="sample_fmt"></param>
        /// <param name="channel_layout"></param>
        /// <param name="sample_rate"></param>
        /// <param name="nb_samples"></param>
        /// <returns></returns>
        private AVFrame* alloc_audio_frame(AVSampleFormat sample_fmt, ulong channel_layout, int sample_rate, int nb_samples)
        {
            AVFrame* frame = ffmpeg.av_frame_alloc();
            if (frame != null)
            {
                frame->format = (int)sample_fmt;
                frame->channel_layout = channel_layout;
                frame->sample_rate = sample_rate;
                frame->nb_samples = nb_samples;
                if (nb_samples != 0)
                {
                    ffmpeg.av_frame_get_buffer(frame, 0).ThrowExceptionIfError();
                }
            }
            return frame;
        }

        /// <summary>
        /// 初始化分配图像帧
        /// </summary>
        /// <param name="pix_fmt"></param>
        /// <param name="width"></param>
        /// <param name="height"></param>
        /// <returns></returns>
        private AVFrame* alloc_video_frame(AVPixelFormat pix_fmt, int width, int height)
        {
            AVFrame* frame;
            frame = ffmpeg.av_frame_alloc();
            if (frame != null)
            {
                frame->format = (int)pix_fmt;
                frame->width = width;
                frame->height = height;
                ffmpeg.av_frame_get_buffer(frame, 0).ThrowExceptionIfError();
            }
            return frame;
        }
        #endregion
    }
}
